{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "from SegTracker import SegTracker\n",
    "from model_args import aot_args,sam_args,segtracker_args\n",
    "from PIL import Image\n",
    "from aot_tracker import _palette\n",
    "import numpy as np\n",
    "import torch\n",
    "import imageio\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.ndimage import binary_dilation\n",
    "import gc\n",
    "\n",
    "\n",
    "\n",
    "def save_prediction(pred_mask,output_dir,file_name):\n",
    "    save_mask = Image.fromarray(pred_mask.astype(np.uint8))\n",
    "    save_mask = save_mask.convert(mode='P')\n",
    "    save_mask.putpalette(_palette)\n",
    "    save_mask.save(os.path.join(output_dir,file_name))\n",
    "def colorize_mask(pred_mask):\n",
    "    save_mask = Image.fromarray(pred_mask.astype(np.uint8))\n",
    "    save_mask = save_mask.convert(mode='P')\n",
    "    save_mask.putpalette(_palette)\n",
    "    save_mask = save_mask.convert(mode='RGB')\n",
    "    return np.array(save_mask)\n",
    "def draw_mask(img, mask, alpha=0.7, id_countour=False):\n",
    "    img_mask = np.zeros_like(img)\n",
    "    img_mask = img\n",
    "    if id_countour:\n",
    "        # very slow ~ 1s per image\n",
    "        obj_ids = np.unique(mask)\n",
    "        obj_ids = obj_ids[obj_ids!=0]\n",
    "\n",
    "        for id in obj_ids:\n",
    "            # Overlay color on  binary mask\n",
    "            if id <= 255:\n",
    "                color = _palette[id*3:id*3+3]\n",
    "            else:\n",
    "                color = [0,0,0]\n",
    "            foreground = img * (1-alpha) + np.ones_like(img) * alpha * np.array(color)\n",
    "            binary_mask = (mask == id)\n",
    "\n",
    "            # Compose image\n",
    "            img_mask[binary_mask] = foreground[binary_mask]\n",
    "\n",
    "            countours = binary_dilation(binary_mask,iterations=1) ^ binary_mask\n",
    "            img_mask[countours, :] = 0\n",
    "    else:\n",
    "        binary_mask = (mask!=0)\n",
    "        countours = binary_dilation(binary_mask,iterations=1) ^ binary_mask\n",
    "        foreground = img*(1-alpha)+colorize_mask(mask)*alpha\n",
    "        img_mask[binary_mask] = foreground[binary_mask]\n",
    "        img_mask[countours,:] = 0\n",
    "        \n",
    "    return img_mask.astype(img.dtype)\n",
    "\n",
    "def create_directories(path):    \n",
    "    dir_path = os.path.dirname(path)\n",
    "    if not os.path.exists(dir_path):\n",
    "        os.makedirs(dir_path)\n",
    "        print(f\"Directory created: {dir_path}\")\n",
    "    else:\n",
    "        print(f\"Directory already exists: {dir_path}\")\n",
    "\n",
    "    return path\n",
    "\n",
    "def video_detection(io_args, segtracker_args, sam_args, aot_args, grounding_caption, box_threshold, text_threshold, box_size_threshold, reset_image):\n",
    "    # source video to segment\n",
    "    import shutil\n",
    "    shutil.copy(io_args['input_video'], io_args['original_video'])\n",
    "    \n",
    "    cap = cv2.VideoCapture(io_args['input_video'])\n",
    "    fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "    # output masks\n",
    "    output_dir = io_args['output_mask_dir']\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.makedirs(output_dir)\n",
    "    pred_list = []\n",
    "    masked_pred_list = []\n",
    "    det_count_one = []\n",
    "    det_count_frames = []\n",
    "    \n",
    "    torch.cuda.empty_cache()\n",
    "    gc.collect()\n",
    "    sam_gap = segtracker_args['sam_gap']\n",
    "    frame_idx = 0\n",
    "    segtracker = SegTracker(segtracker_args, sam_args, aot_args)\n",
    "    segtracker.restart_tracker()\n",
    "    \n",
    "    with torch.cuda.amp.autocast():\n",
    "        while cap.isOpened():\n",
    "            ret, frame = cap.read()\n",
    "            if not ret:\n",
    "                break\n",
    "            frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
    "            if frame_idx == 0:\n",
    "                pred_mask, _ = segtracker.detect_and_seg(frame, grounding_caption, box_threshold, text_threshold, box_size_threshold, reset_image)\n",
    "                torch.cuda.empty_cache()\n",
    "                gc.collect()\n",
    "                segtracker.add_reference(frame, pred_mask)\n",
    "            elif (frame_idx % sam_gap) == 0:\n",
    "                seg_mask, _ = segtracker.detect_and_seg(frame, grounding_caption, box_threshold, text_threshold, box_size_threshold, reset_image)\n",
    "                save_prediction(seg_mask, './debug/seg_result', str(frame_idx)+'.png')\n",
    "                torch.cuda.empty_cache()\n",
    "                gc.collect()\n",
    "                track_mask = segtracker.track(frame)\n",
    "                save_prediction(track_mask, './debug/aot_result', str(frame_idx)+'.png')\n",
    "                new_obj_mask = segtracker.find_new_objs(track_mask, seg_mask)\n",
    "                if np.sum(new_obj_mask > 0) >  frame.shape[0] * frame.shape[1] * 0.4:\n",
    "                    new_obj_mask = np.zeros_like(new_obj_mask)\n",
    "                save_prediction(new_obj_mask,output_dir,str(frame_idx)+'_new.png')\n",
    "                pred_mask = track_mask + new_obj_mask\n",
    "                segtracker.add_reference(frame, pred_mask)\n",
    "            else:\n",
    "                pred_mask = segtracker.track(frame,update_memory=True)\n",
    "            torch.cuda.empty_cache()\n",
    "            gc.collect()\n",
    "            \n",
    "            save_prediction(pred_mask,output_dir,str(frame_idx)+'.png')\n",
    "            pred_list.append(pred_mask)\n",
    "\n",
    "            obj_ids = np.unique(pred_mask)\n",
    "            obj_ids = obj_ids[obj_ids!=0]\n",
    "            det_count_frames.append(len(obj_ids))\n",
    "\n",
    "            print(\"processed frame {}, obj_num {}\".format(frame_idx,segtracker.get_obj_num()),end='\\r')\n",
    "               \n",
    "            frame_idx += 1\n",
    "        # # det_count_one = np.sum(det_count_frames)/frame_idx\n",
    "        # det_count_avg = np.sum(det_count_frames) / frame_idx\n",
    "        # det_count_one = np.sum(np.abs(det_count_frames - det_count_avg))/frame_idx\n",
    "        # cap.release()\n",
    "        \n",
    "        # draw pred mask on frame and save as a video\n",
    "        cap = cv2.VideoCapture(io_args['input_video'])\n",
    "        fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "        num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "        \n",
    "        if io_args['input_video'][-3:]=='mp4':\n",
    "            fourcc =  cv2.VideoWriter_fourcc(*\"mp4v\")\n",
    "        elif io_args['input_video'][-3:] == 'avi':\n",
    "            fourcc =  cv2.VideoWriter_fourcc(*\"MJPG\")\n",
    "            # fourcc = cv2.VideoWriter_fourcc(*\"XVID\")\n",
    "        else:\n",
    "            fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))\n",
    "        out = cv2.VideoWriter(io_args['output_video'], fourcc, fps, (width, height))\n",
    "        frame_idx = 0\n",
    "        while cap.isOpened():\n",
    "            ret, frame = cap.read()\n",
    "            if not ret:\n",
    "                break\n",
    "            frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
    "            pred_mask = pred_list[frame_idx]\n",
    "            masked_frame = draw_mask(frame,pred_mask)\n",
    "            # masked_frame = masked_pred_list[frame_idx]\n",
    "            masked_frame = cv2.cvtColor(masked_frame,cv2.COLOR_RGB2BGR)\n",
    "            out.write(masked_frame)\n",
    "            print('frame {} writed'.format(frame_idx),end='\\r')\n",
    "            frame_idx += 1\n",
    "        out.release()\n",
    "        cap.release()\n",
    "        print(\"\\n{} saved\".format(io_args['output_video']))\n",
    "        # save colorized masks as a gif\n",
    "        duration = 1000 / fps\n",
    "        imageio.mimsave(io_args['output_gif'],pred_list,duration=duration)\n",
    "        print(\"{} saved\".format(io_args['output_gif']))\n",
    "        print('\\nfinished')\n",
    "        \n",
    "    return pred_list, det_count_frames \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(video_paths[1])\n",
    "# print(prompt_paths[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conduct evaluation using COCO classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "Finishing last run (ID:6u88w020) before initializing another..."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">20230825-180730_pikavideos</strong> at: <a href='https://wandb.ai/bruceliu/Vid%20Model%20Eval/runs/6u88w020' target=\"_blank\">https://wandb.ai/bruceliu/Vid%20Model%20Eval/runs/6u88w020</a><br/> View job at <a href='https://wandb.ai/bruceliu/Vid%20Model%20Eval/jobs/QXJ0aWZhY3RDb2xsZWN0aW9uOjkyMzQ3NzE0/version_details/v3' target=\"_blank\">https://wandb.ai/bruceliu/Vid%20Model%20Eval/jobs/QXJ0aWZhY3RDb2xsZWN0aW9uOjkyMzQ3NzE0/version_details/v3</a><br/>Synced 5 W&B file(s), 0 media file(s), 2 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20230825_180736-6u88w020/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Successfully finished last run (ID:6u88w020). Initializing new run:<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "94d908a904c84b3a862adbacaf8f2e3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.01666861826670356, max=1.0)…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.8"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/apdcephfs_cq3/share_1290939/raphaelliu/Segment-and-Track-Anything/wandb/run-20230825_180853-cdoi55tp</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/bruceliu/Vid%20Model%20Eval/runs/cdoi55tp' target=\"_blank\">20230825-180853_pikavideos</a></strong> to <a href='https://wandb.ai/bruceliu/Vid%20Model%20Eval' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/bruceliu/Vid%20Model%20Eval' target=\"_blank\">https://wandb.ai/bruceliu/Vid%20Model%20Eval</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/bruceliu/Vid%20Model%20Eval/runs/cdoi55tp' target=\"_blank\">https://wandb.ai/bruceliu/Vid%20Model%20Eval/runs/cdoi55tp</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:09:07 Total video num 53622\n",
      "2023-08-25 18:09:07 Total video num 53622\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Directory created: ./pikavideos/20230825-180853/bird/masks\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/STA/lib/python3.9/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3483.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/STA/lib/python3.9/site-packages/transformers/modeling_utils.py:881: FutureWarning: The `device` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n",
      "/root/miniconda3/envs/STA/lib/python3.9/site-packages/torch/utils/checkpoint.py:31: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
      "  warnings.warn(\"None of the inputs have requires_grad=True. Gradients will be None\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 23 writed 23, obj_num 0\n",
      "./pikavideos/20230825-180853/bird/mask_Imagine_stepping_into_a_vast_and_enchanting_forest._Towering_trees_stretch_upwards_their_branches_f_seed6872407215330904296.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:09:20 Det num: 1, Positive num: 0, Acc: 0.0, Step obj-gen score: 0.0, Whole obj-gen score: 1.0, Vid num: 1, Prompt:  Imagine stepping into a vast and enchanting forest. Towering trees stretch upwards, their branches forming a lush canopy that filters sunlight and creates dancing patterns on the forest floor. The air is filled with the earthy aroma of damp moss, pine needles, and the subtle scent of wildflowers.  As you walk deeper into the woods, a sense of serenity envelops you. The ground is blanketed with a rich layer of fallen leaves and soft undergrowth, muffling your footsteps. The rustling of leaves and the distant call of birds create a soothing symphony of nature.  Shafts of sunlight pierce through the gaps in the canopy, illuminating patches of vibrant green ferns and delicate wildflowers. Mushrooms of various shapes and colors decorate the base of ancient trees, adding a touch of whimsy to the scene.  A small stream trickles nearby, its clear water bubbling over smooth stones. The tranquil sound of running water provides a gentle soundtrack to your exploration. As you continue your journey, you might spot the occasional glimpse of wildlife â a squirrel darting up a tree, a rabbit nibbling on leaves, or perhaps even a deer cautiously watching from a distance.  Author: @victor, Path index: 11, Video name: Imagine_stepping_into_a_vast_and_enchanting_forest._Towering_trees_stretch_upwards_their_branches_f_seed6872407215330904296\n",
      "2023-08-25 18:09:20 Det num: 1, Positive num: 0, Acc: 0.0, Step obj-gen score: 0.0, Whole obj-gen score: 1.0, Vid num: 1, Prompt:  Imagine stepping into a vast and enchanting forest. Towering trees stretch upwards, their branches forming a lush canopy that filters sunlight and creates dancing patterns on the forest floor. The air is filled with the earthy aroma of damp moss, pine needles, and the subtle scent of wildflowers.  As you walk deeper into the woods, a sense of serenity envelops you. The ground is blanketed with a rich layer of fallen leaves and soft undergrowth, muffling your footsteps. The rustling of leaves and the distant call of birds create a soothing symphony of nature.  Shafts of sunlight pierce through the gaps in the canopy, illuminating patches of vibrant green ferns and delicate wildflowers. Mushrooms of various shapes and colors decorate the base of ancient trees, adding a touch of whimsy to the scene.  A small stream trickles nearby, its clear water bubbling over smooth stones. The tranquil sound of running water provides a gentle soundtrack to your exploration. As you continue your journey, you might spot the occasional glimpse of wildlife â a squirrel darting up a tree, a rabbit nibbling on leaves, or perhaps even a deer cautiously watching from a distance.  Author: @victor, Path index: 11, Video name: Imagine_stepping_into_a_vast_and_enchanting_forest._Towering_trees_stretch_upwards_their_branches_f_seed6872407215330904296\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/bird/gif_Imagine_stepping_into_a_vast_and_enchanting_forest._Towering_trees_stretch_upwards_their_branches_f_seed6872407215330904296.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/elephant/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 1\n",
      "./pikavideos/20230825-180853/elephant/mask_Naruto_chasing_an_Elephant_in_Africa_anime_Naruto-punk_cute_cartoon_fast_action_seed12455689736558127025.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:09:31 Det num: 2, Positive num: 1, Acc: 0.5, Step obj-gen score: 0.0, Whole obj-gen score: 0.5, Vid num: 2, Prompt:  Naruto chasing an Elephant in Africa, anime Naruto-punk cute cartoon fast action  Author: @RubenTainoAI, Path index: 12, Video name: Naruto_chasing_an_Elephant_in_Africa_anime_Naruto-punk_cute_cartoon_fast_action_seed12455689736558127025\n",
      "2023-08-25 18:09:31 Det num: 2, Positive num: 1, Acc: 0.5, Step obj-gen score: 0.0, Whole obj-gen score: 0.5, Vid num: 2, Prompt:  Naruto chasing an Elephant in Africa, anime Naruto-punk cute cartoon fast action  Author: @RubenTainoAI, Path index: 12, Video name: Naruto_chasing_an_Elephant_in_Africa_anime_Naruto-punk_cute_cartoon_fast_action_seed12455689736558127025\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/elephant/gif_Naruto_chasing_an_Elephant_in_Africa_anime_Naruto-punk_cute_cartoon_fast_action_seed12455689736558127025.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/dog/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 1\n",
      "./pikavideos/20230825-180853/dog/mask_attractive_angry_female_zombie_wearing_dog_collar_and_chain_running_trying_to_attack_camera_fierce_seed14017085236677006407.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:09:41 Det num: 3, Positive num: 2, Acc: 0.6666666666666666, Step obj-gen score: 0.0, Whole obj-gen score: 0.3333333333333333, Vid num: 3, Prompt:  attractive angry female zombie wearing dog collar and chain, running, trying to attack camera fiercely and with determination. Ultra photo realistic 8k resolution --ar 16:9  Author: @IzzyAi, Path index: 14, Video name: attractive_angry_female_zombie_wearing_dog_collar_and_chain_running_trying_to_attack_camera_fierce_seed14017085236677006407\n",
      "2023-08-25 18:09:41 Det num: 3, Positive num: 2, Acc: 0.6666666666666666, Step obj-gen score: 0.0, Whole obj-gen score: 0.3333333333333333, Vid num: 3, Prompt:  attractive angry female zombie wearing dog collar and chain, running, trying to attack camera fiercely and with determination. Ultra photo realistic 8k resolution --ar 16:9  Author: @IzzyAi, Path index: 14, Video name: attractive_angry_female_zombie_wearing_dog_collar_and_chain_running_trying_to_attack_camera_fierce_seed14017085236677006407\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/dog/gif_attractive_angry_female_zombie_wearing_dog_collar_and_chain_running_trying_to_attack_camera_fierce_seed14017085236677006407.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/horse/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 1\n",
      "./pikavideos/20230825-180853/horse/mask_a_red_horse_with_robotic_legs_and_eyes_running_alongside_the_beach_-_ar_9_16_seed18364374028985219773.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:09:51 Det num: 4, Positive num: 3, Acc: 0.75, Step obj-gen score: 0.0, Whole obj-gen score: 0.25, Vid num: 4, Prompt:  a red horse with robotic legs and eyes, running alongside the beach - ar 9:16  Author: @imgcatz, Path index: 38, Video name: a_red_horse_with_robotic_legs_and_eyes_running_alongside_the_beach_-_ar_9_16_seed18364374028985219773\n",
      "2023-08-25 18:09:51 Det num: 4, Positive num: 3, Acc: 0.75, Step obj-gen score: 0.0, Whole obj-gen score: 0.25, Vid num: 4, Prompt:  a red horse with robotic legs and eyes, running alongside the beach - ar 9:16  Author: @imgcatz, Path index: 38, Video name: a_red_horse_with_robotic_legs_and_eyes_running_alongside_the_beach_-_ar_9_16_seed18364374028985219773\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/horse/gif_a_red_horse_with_robotic_legs_and_eyes_running_alongside_the_beach_-_ar_9_16_seed18364374028985219773.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/bear/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 0\n",
      "./pikavideos/20230825-180853/bear/mask_MY_TESTAMENT_When_I_am_dead_bury_me_In_my_beloved_Ukraine_My_tomb_upon_a_grave_mound_high_Amid_the_seed3133089823467916576.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:10:01 Det num: 5, Positive num: 3, Acc: 0.6, Step obj-gen score: 0.0, Whole obj-gen score: 0.2, Vid num: 5, Prompt:  MY TESTAMENT When I am dead, bury me In my beloved Ukraine, My tomb upon a grave mound high Amid the spreading plain, So that the fields, the boundless steppes, The Dnipro's plunging shore My eyes could see, my ears could hear The mighty river roar. When from Ukraine the Dnipro bears Into the deep blue sea The blood of foes ... then will I leave These hills and fertile fields -- I'll leave them all and fly away To the abode of God, And then I'll pray ....  But till that dayI nothing know of God. bury me, then rise ye up And break your heavy chains And water with the tyrants' blood The freedom you have gained. And in the great new family, The family of the free, With softly spoken, kindly word Remember also me.  Author: @Ashtray25, Path index: 42, Video name: MY_TESTAMENT_When_I_am_dead_bury_me_In_my_beloved_Ukraine_My_tomb_upon_a_grave_mound_high_Amid_the_seed3133089823467916576\n",
      "2023-08-25 18:10:01 Det num: 5, Positive num: 3, Acc: 0.6, Step obj-gen score: 0.0, Whole obj-gen score: 0.2, Vid num: 5, Prompt:  MY TESTAMENT When I am dead, bury me In my beloved Ukraine, My tomb upon a grave mound high Amid the spreading plain, So that the fields, the boundless steppes, The Dnipro's plunging shore My eyes could see, my ears could hear The mighty river roar. When from Ukraine the Dnipro bears Into the deep blue sea The blood of foes ... then will I leave These hills and fertile fields -- I'll leave them all and fly away To the abode of God, And then I'll pray ....  But till that dayI nothing know of God. bury me, then rise ye up And break your heavy chains And water with the tyrants' blood The freedom you have gained. And in the great new family, The family of the free, With softly spoken, kindly word Remember also me.  Author: @Ashtray25, Path index: 42, Video name: MY_TESTAMENT_When_I_am_dead_bury_me_In_my_beloved_Ukraine_My_tomb_upon_a_grave_mound_high_Amid_the_seed3133089823467916576\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/bear/gif_MY_TESTAMENT_When_I_am_dead_bury_me_In_my_beloved_Ukraine_My_tomb_upon_a_grave_mound_high_Amid_the_seed3133089823467916576.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/cat/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 21 writed 23, obj_num 1\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:10:11 Det num: 6, Positive num: 4, Acc: 0.6666666666666666, Step obj-gen score: 0.0, Whole obj-gen score: 0.16666666666666666, Vid num: 6, Prompt:  Prompt: Cinematic Composition of a cat walking in a space  Author: @mksparks, Path index: 49, Video name: Prompt__Cinematic_Composition_of_a_cat_walking_in_a_space_seed7257964282371214865\n",
      "2023-08-25 18:10:11 Det num: 6, Positive num: 4, Acc: 0.6666666666666666, Step obj-gen score: 0.0, Whole obj-gen score: 0.16666666666666666, Vid num: 6, Prompt:  Prompt: Cinematic Composition of a cat walking in a space  Author: @mksparks, Path index: 49, Video name: Prompt__Cinematic_Composition_of_a_cat_walking_in_a_space_seed7257964282371214865\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 23 writed\n",
      "./pikavideos/20230825-180853/cat/mask_Prompt__Cinematic_Composition_of_a_cat_walking_in_a_space_seed7257964282371214865.mp4 saved\n",
      "./pikavideos/20230825-180853/cat/gif_Prompt__Cinematic_Composition_of_a_cat_walking_in_a_space_seed7257964282371214865.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/toilet/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 0\n",
      "./pikavideos/20230825-180853/toilet/mask_war_scene_with_toilets_and_mutatant_camera_head_people_with_explosions_and_broken_toilets_on_fire_wi_seed9022266552524930572.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:10:21 Det num: 7, Positive num: 4, Acc: 0.5714285714285714, Step obj-gen score: 0.0, Whole obj-gen score: 0.14285714285714285, Vid num: 7, Prompt:  war scene with toilets and mutatant camera head people with explosions and broken toilets on fire with explosions in the background and giant flying cameras  Author: @astralmoonveil, Path index: 57, Video name: war_scene_with_toilets_and_mutatant_camera_head_people_with_explosions_and_broken_toilets_on_fire_wi_seed9022266552524930572\n",
      "2023-08-25 18:10:21 Det num: 7, Positive num: 4, Acc: 0.5714285714285714, Step obj-gen score: 0.0, Whole obj-gen score: 0.14285714285714285, Vid num: 7, Prompt:  war scene with toilets and mutatant camera head people with explosions and broken toilets on fire with explosions in the background and giant flying cameras  Author: @astralmoonveil, Path index: 57, Video name: war_scene_with_toilets_and_mutatant_camera_head_people_with_explosions_and_broken_toilets_on_fire_wi_seed9022266552524930572\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/toilet/gif_war_scene_with_toilets_and_mutatant_camera_head_people_with_explosions_and_broken_toilets_on_fire_wi_seed9022266552524930572.gif saved\n",
      "\n",
      "finished\n",
      "Directory already exists: ./pikavideos/20230825-180853/dog/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 21 writed 23, obj_num 2\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:10:31 Det num: 8, Positive num: 5, Acc: 0.625, Step obj-gen score: 0.0, Whole obj-gen score: 0.125, Vid num: 8, Prompt: @Robi A Swiss Mountain Dog sits gracefully in front of the Grand Canyon. The sun shines brightly, and a gentle breeze 4K -ar 16:9  Author: @Robi@Kali, Path index: 64, Video name: A_Swiss_Mountain_Dog_sits_gracefully_in_front_of_the_Grand_Canyon._The_sun_shines_brightly_and_a_ge_seed8816516681728766960\n",
      "2023-08-25 18:10:31 Det num: 8, Positive num: 5, Acc: 0.625, Step obj-gen score: 0.0, Whole obj-gen score: 0.125, Vid num: 8, Prompt: @Robi A Swiss Mountain Dog sits gracefully in front of the Grand Canyon. The sun shines brightly, and a gentle breeze 4K -ar 16:9  Author: @Robi@Kali, Path index: 64, Video name: A_Swiss_Mountain_Dog_sits_gracefully_in_front_of_the_Grand_Canyon._The_sun_shines_brightly_and_a_ge_seed8816516681728766960\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 23 writed\n",
      "./pikavideos/20230825-180853/dog/mask_A_Swiss_Mountain_Dog_sits_gracefully_in_front_of_the_Grand_Canyon._The_sun_shines_brightly_and_a_ge_seed8816516681728766960.mp4 saved\n",
      "./pikavideos/20230825-180853/dog/gif_A_Swiss_Mountain_Dog_sits_gracefully_in_front_of_the_Grand_Canyon._The_sun_shines_brightly_and_a_ge_seed8816516681728766960.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/cow/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 20 writed 23, obj_num 0\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:10:41 Det num: 9, Positive num: 5, Acc: 0.5555555555555556, Step obj-gen score: 0.0, Whole obj-gen score: 0.1111111111111111, Vid num: 9, Prompt:  a video of President Bolsoaro in the midst of several cows, super realistic --ar 16:9  Author: @INTERNACIONA, Path index: 79, Video name: a_video_of_President_Bolsoaro_in_the_midst_of_several_cows_super_realistic_--ar_16_9\n",
      "2023-08-25 18:10:41 Det num: 9, Positive num: 5, Acc: 0.5555555555555556, Step obj-gen score: 0.0, Whole obj-gen score: 0.1111111111111111, Vid num: 9, Prompt:  a video of President Bolsoaro in the midst of several cows, super realistic --ar 16:9  Author: @INTERNACIONA, Path index: 79, Video name: a_video_of_President_Bolsoaro_in_the_midst_of_several_cows_super_realistic_--ar_16_9\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 23 writed\n",
      "./pikavideos/20230825-180853/cow/mask_a_video_of_President_Bolsoaro_in_the_midst_of_several_cows_super_realistic_--ar_16_9.mp4 saved\n",
      "./pikavideos/20230825-180853/cow/gif_a_video_of_President_Bolsoaro_in_the_midst_of_several_cows_super_realistic_--ar_16_9.gif saved\n",
      "\n",
      "finished\n",
      "Directory already exists: ./pikavideos/20230825-180853/dog/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 1\n",
      "./pikavideos/20230825-180853/dog/mask_a_white_Canaan_dog_flying_with_a_red_cape_seed4449926014351470217.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:10:51 Det num: 10, Positive num: 6, Acc: 0.6, Step obj-gen score: 0.0, Whole obj-gen score: 0.1, Vid num: 10, Prompt:  a white Canaan dog flying with a red cape  Author: @poeticoncept, Path index: 86, Video name: a_white_Canaan_dog_flying_with_a_red_cape_seed4449926014351470217\n",
      "2023-08-25 18:10:51 Det num: 10, Positive num: 6, Acc: 0.6, Step obj-gen score: 0.0, Whole obj-gen score: 0.1, Vid num: 10, Prompt:  a white Canaan dog flying with a red cape  Author: @poeticoncept, Path index: 86, Video name: a_white_Canaan_dog_flying_with_a_red_cape_seed4449926014351470217\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/dog/gif_a_white_Canaan_dog_flying_with_a_red_cape_seed4449926014351470217.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/cake/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 1\n",
      "./pikavideos/20230825-180853/cake/mask_cinematic_black_and_white_noir_close_up_film_of_frog_chefs_throwing_cake_at_each_other_exploding_ca_seed17018909320981396447.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:11:01 Det num: 11, Positive num: 7, Acc: 0.6363636363636364, Step obj-gen score: 0.0, Whole obj-gen score: 0.09090909090909091, Vid num: 11, Prompt:  cinematic black and white noir close up film of frog chefs throwing cake at each other, exploding cakes, surreal, weird, twisted, retro, 8k  Author: @plannetina, Path index: 87, Video name: cinematic_black_and_white_noir_close_up_film_of_frog_chefs_throwing_cake_at_each_other_exploding_ca_seed17018909320981396447\n",
      "2023-08-25 18:11:01 Det num: 11, Positive num: 7, Acc: 0.6363636363636364, Step obj-gen score: 0.0, Whole obj-gen score: 0.09090909090909091, Vid num: 11, Prompt:  cinematic black and white noir close up film of frog chefs throwing cake at each other, exploding cakes, surreal, weird, twisted, retro, 8k  Author: @plannetina, Path index: 87, Video name: cinematic_black_and_white_noir_close_up_film_of_frog_chefs_throwing_cake_at_each_other_exploding_ca_seed17018909320981396447\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/cake/gif_cinematic_black_and_white_noir_close_up_film_of_frog_chefs_throwing_cake_at_each_other_exploding_ca_seed17018909320981396447.gif saved\n",
      "\n",
      "finished\n",
      "Directory already exists: ./pikavideos/20230825-180853/bear/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 0\n",
      "./pikavideos/20230825-180853/bear/mask_bear_breakdancing_in_the_trading_room_in_front_of_the_screens_depicting_candlestick_charts_red_going.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:11:11 Det num: 12, Positive num: 7, Acc: 0.5833333333333334, Step obj-gen score: 0.0, Whole obj-gen score: 0.08333333333333333, Vid num: 12, Prompt:  bear breakdancing in the trading room in front of the screens depicting candlestick charts red going down -neg cartoon, green -sg 99  Author: @MickeyMoloch, Path index: 103, Video name: bear_breakdancing_in_the_trading_room_in_front_of_the_screens_depicting_candlestick_charts_red_going\n",
      "2023-08-25 18:11:11 Det num: 12, Positive num: 7, Acc: 0.5833333333333334, Step obj-gen score: 0.0, Whole obj-gen score: 0.08333333333333333, Vid num: 12, Prompt:  bear breakdancing in the trading room in front of the screens depicting candlestick charts red going down -neg cartoon, green -sg 99  Author: @MickeyMoloch, Path index: 103, Video name: bear_breakdancing_in_the_trading_room_in_front_of_the_screens_depicting_candlestick_charts_red_going\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/bear/gif_bear_breakdancing_in_the_trading_room_in_front_of_the_screens_depicting_candlestick_charts_red_going.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/book/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 0\n",
      "./pikavideos/20230825-180853/book/mask_A_student_sitting_at_a_desk_with_a_book_showing_a_determined_expression_while_studying.hype_realist_seed10023452760357934235.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:11:21 Det num: 13, Positive num: 7, Acc: 0.5384615384615384, Step obj-gen score: 0.0, Whole obj-gen score: 0.07692307692307693, Vid num: 13, Prompt:  A student sitting at a desk with a book, showing a determined expression while studying.hype realistic, 8k  Author: @Raihan, Path index: 104, Video name: A_student_sitting_at_a_desk_with_a_book_showing_a_determined_expression_while_studying.hype_realist_seed10023452760357934235\n",
      "2023-08-25 18:11:21 Det num: 13, Positive num: 7, Acc: 0.5384615384615384, Step obj-gen score: 0.0, Whole obj-gen score: 0.07692307692307693, Vid num: 13, Prompt:  A student sitting at a desk with a book, showing a determined expression while studying.hype realistic, 8k  Author: @Raihan, Path index: 104, Video name: A_student_sitting_at_a_desk_with_a_book_showing_a_determined_expression_while_studying.hype_realist_seed10023452760357934235\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/book/gif_A_student_sitting_at_a_desk_with_a_book_showing_a_determined_expression_while_studying.hype_realist_seed10023452760357934235.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/boat/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 21 writed 23, obj_num 1\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:11:31 Det num: 14, Positive num: 8, Acc: 0.5714285714285714, Step obj-gen score: 0.0, Whole obj-gen score: 0.07142857142857142, Vid num: 14, Prompt:  huge waves crash the boat, 4k, storm, epic, cinematic  Author: @Mike, Path index: 113, Video name: huge_waves_crash_the_boat_4k_storm_epic_cinematic_seed1568166036430126724\n",
      "2023-08-25 18:11:31 Det num: 14, Positive num: 8, Acc: 0.5714285714285714, Step obj-gen score: 0.0, Whole obj-gen score: 0.07142857142857142, Vid num: 14, Prompt:  huge waves crash the boat, 4k, storm, epic, cinematic  Author: @Mike, Path index: 113, Video name: huge_waves_crash_the_boat_4k_storm_epic_cinematic_seed1568166036430126724\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 23 writed\n",
      "./pikavideos/20230825-180853/boat/mask_huge_waves_crash_the_boat_4k_storm_epic_cinematic_seed1568166036430126724.mp4 saved\n",
      "./pikavideos/20230825-180853/boat/gif_huge_waves_crash_the_boat_4k_storm_epic_cinematic_seed1568166036430126724.gif saved\n",
      "\n",
      "finished\n",
      "Directory created: ./pikavideos/20230825-180853/car/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 2\n",
      "./pikavideos/20230825-180853/car/mask_super_cars_and_superbikes_racing_in_konkan_slow_motion_hyper_realistic_4k_resolution_high_octane_seed4192569648458836432.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:11:42 Det num: 15, Positive num: 9, Acc: 0.6, Step obj-gen score: 0.9166666666666664, Whole obj-gen score: 0.005555555555555573, Vid num: 15, Prompt:  super cars and superbikes racing in konkan, slow motion, hyper realistic, 4k resolution, high octane render, golden hour lighting, intelligent render, beyond human perspective, dynamic shot, blade runner 2049 style  Author: @rohitmokashi90, Path index: 121, Video name: super_cars_and_superbikes_racing_in_konkan_slow_motion_hyper_realistic_4k_resolution_high_octane_seed4192569648458836432\n",
      "2023-08-25 18:11:42 Det num: 15, Positive num: 9, Acc: 0.6, Step obj-gen score: 0.9166666666666664, Whole obj-gen score: 0.005555555555555573, Vid num: 15, Prompt:  super cars and superbikes racing in konkan, slow motion, hyper realistic, 4k resolution, high octane render, golden hour lighting, intelligent render, beyond human perspective, dynamic shot, blade runner 2049 style  Author: @rohitmokashi90, Path index: 121, Video name: super_cars_and_superbikes_racing_in_konkan_slow_motion_hyper_realistic_4k_resolution_high_octane_seed4192569648458836432\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/car/gif_super_cars_and_superbikes_racing_in_konkan_slow_motion_hyper_realistic_4k_resolution_high_octane_seed4192569648458836432.gif saved\n",
      "\n",
      "finished\n",
      "Directory already exists: ./pikavideos/20230825-180853/elephant/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 2\n",
      "./pikavideos/20230825-180853/elephant/mask_African_Elephant_nature_photography_expansive_awe-inspiring_breathtaking_vivid_colors_dramatic__seed6798599620705850248.mp4 saved\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:11:52 Det num: 16, Positive num: 10, Acc: 0.625, Step obj-gen score: 0.0, Whole obj-gen score: 0.0052083333333333495, Vid num: 16, Prompt:  African Elephant nature photography, expansive, awe-inspiring, breathtaking, vivid colors, dramatic lighting, wide-angle, sharp focus, good exposure, golden hour -ar 16:9  Author: @efadiman, Path index: 127, Video name: African_Elephant_nature_photography_expansive_awe-inspiring_breathtaking_vivid_colors_dramatic__seed6798599620705850248\n",
      "2023-08-25 18:11:52 Det num: 16, Positive num: 10, Acc: 0.625, Step obj-gen score: 0.0, Whole obj-gen score: 0.0052083333333333495, Vid num: 16, Prompt:  African Elephant nature photography, expansive, awe-inspiring, breathtaking, vivid colors, dramatic lighting, wide-angle, sharp focus, good exposure, golden hour -ar 16:9  Author: @efadiman, Path index: 127, Video name: African_Elephant_nature_photography_expansive_awe-inspiring_breathtaking_vivid_colors_dramatic__seed6798599620705850248\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./pikavideos/20230825-180853/elephant/gif_African_Elephant_nature_photography_expansive_awe-inspiring_breathtaking_vivid_colors_dramatic__seed6798599620705850248.gif saved\n",
      "\n",
      "finished\n",
      "Directory already exists: ./pikavideos/20230825-180853/horse/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 21 writed 23, obj_num 1\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:12:02 Det num: 17, Positive num: 11, Acc: 0.6470588235294118, Step obj-gen score: 0.0, Whole obj-gen score: 0.004901960784313741, Vid num: 17, Prompt:  two cowboys riding horses near a farm during American Civil War  Author: @fahadqureashi, Path index: 136, Video name: two_cowboys_riding_horses_near_a_farm_during_American_Civil_War_seed5697136767728898480\n",
      "2023-08-25 18:12:02 Det num: 17, Positive num: 11, Acc: 0.6470588235294118, Step obj-gen score: 0.0, Whole obj-gen score: 0.004901960784313741, Vid num: 17, Prompt:  two cowboys riding horses near a farm during American Civil War  Author: @fahadqureashi, Path index: 136, Video name: two_cowboys_riding_horses_near_a_farm_during_American_Civil_War_seed5697136767728898480\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 23 writed\n",
      "./pikavideos/20230825-180853/horse/mask_two_cowboys_riding_horses_near_a_farm_during_American_Civil_War_seed5697136767728898480.mp4 saved\n",
      "./pikavideos/20230825-180853/horse/gif_two_cowboys_riding_horses_near_a_farm_during_American_Civil_War_seed5697136767728898480.gif saved\n",
      "\n",
      "finished\n",
      "Directory already exists: ./pikavideos/20230825-180853/toilet/masks\n",
      "final text_encoder_type: bert-base-uncased\n",
      "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
      " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
      "SegTracker has been initialized\n",
      "frame 23 writed 23, obj_num 0\n",
      "./pikavideos/20230825-180853/toilet/mask_skibidi_toilet_dancing_seed16561181095301632481.mp4 saved\n",
      "./pikavideos/20230825-180853/toilet/gif_skibidi_toilet_dancing_seed16561181095301632481.gif saved\n",
      "\n",
      "finished\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-25 18:12:12 Det num: 18, Positive num: 11, Acc: 0.6111111111111112, Step obj-gen score: 0.0, Whole obj-gen score: 0.004629629629629644, Vid num: 18, Prompt:  skibidi toilet dancing  Author: @crew, Path index: 137, Video name: skibidi_toilet_dancing_seed16561181095301632481\n",
      "2023-08-25 18:12:12 Det num: 18, Positive num: 11, Acc: 0.6111111111111112, Step obj-gen score: 0.0, Whole obj-gen score: 0.004629629629629644, Vid num: 18, Prompt:  skibidi toilet dancing  Author: @crew, Path index: 137, Video name: skibidi_toilet_dancing_seed16561181095301632481\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Directory already exists: ./pikavideos/20230825-180853/bird/masks\n",
      "final text_encoder_type: bert-base-uncased\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import os\n",
    "import re\n",
    "import logging\n",
    "import time\n",
    "import wandb\n",
    "\n",
    "#### Detection on videos ####\n",
    "# task = 'fulljourney_videos'\n",
    "task = 'pikavideos'\n",
    "# Create the directory if it doesn't exist\n",
    "timestamp = time.strftime(\"%Y%m%d-%H%M%S\")\n",
    "\n",
    "os.makedirs(f\"./{task}/{timestamp}\", exist_ok=True)\n",
    "wandb.init(project=\"Vid Model Eval\",name=str(timestamp)+\"_\"+task)\n",
    "\n",
    "# Set up logging\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "# File handler for writing logs to a file\n",
    "file_handler = logging.FileHandler(filename=f\"./{task}/{timestamp}/object_det_record.txt\")\n",
    "file_handler.setFormatter(logging.Formatter(\"%(asctime)s %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\"))\n",
    "logger.addHandler(file_handler)\n",
    "\n",
    "# Stream handler for displaying logs in the terminal\n",
    "stream_handler = logging.StreamHandler()\n",
    "stream_handler.setFormatter(logging.Formatter(\"%(asctime)s %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\"))\n",
    "logger.addHandler(stream_handler)\n",
    "\n",
    "# settings for detection function\n",
    "# choose good parameters in sam_args based on the first frame segmentation result\n",
    "# other arguments can be modified in model_args.py\n",
    "# note the object number limit is 255 by default, which requires < 10GB GPU memory with amp\n",
    "sam_args['generator_args'] = {\n",
    "        'points_per_side': 30,\n",
    "        'pred_iou_thresh': 0.8,\n",
    "        'stability_score_thresh': 0.9,\n",
    "        'crop_n_layers': 1,\n",
    "        'crop_n_points_downscale_factor': 2,\n",
    "        'min_mask_region_area': 200,\n",
    "    }\n",
    "\n",
    "# For every sam_gap frames, we use SAM to find new objects and add them for tracking\n",
    "# larger sam_gap is faster but may not spot new objects in time\n",
    "segtracker_args = {\n",
    "    'sam_gap': 49, # the interval to run sam to segment new objects\n",
    "    'min_area': 200, # minimal mask area to add a new mask as a new object\n",
    "    'max_obj_num': 255, # maximal object number to track in a video\n",
    "    'min_new_obj_iou': 0.8, # the area of a new object in the background should > 80% \n",
    "}\n",
    "\n",
    "# Set Text args\n",
    "'''\n",
    "parameter:\n",
    "    grounding_caption: Text prompt to detect objects in key-frames\n",
    "    box_threshold: threshold for box \n",
    "    text_threshold: threshold for label(text)\n",
    "    box_size_threshold: If the size ratio between the box and the frame is larger than the box_size_threshold, the box will be ignored. This is used to filter out large boxes.\n",
    "    reset_image: reset the image embeddings for SAM\n",
    "'''\n",
    "# grounding_caption = \"car.suv\"\n",
    "grounding_caption = \"\" #must have this class in the image, otherwise go wrong\n",
    "box_threshold, text_threshold, box_size_threshold, reset_image = 0.6, 0.5, 0.5, True \n",
    "\n",
    "# COCO dataset\n",
    "keywords = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',\n",
    "                   'train', 'truck', 'boat', 'traffic light', 'fire hydrant',\n",
    "                   'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',\n",
    "                   'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',\n",
    "                   'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',\n",
    "                   'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',\n",
    "                   'baseball glove', 'skateboard', 'surfboard', 'tennis racket',\n",
    "                   'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',\n",
    "                   'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',\n",
    "                   'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',\n",
    "                   'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',\n",
    "                   'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',\n",
    "                   'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',\n",
    "                   'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] #coco classes, https://github.com/matlab-deep-learning/Object-Detection-Using-Pretrained-YOLO-v2/blob/main/+helper/coco-classes.txt\n",
    "    \n",
    "    \n",
    "dir_videos = f'/apdcephfs_cq3/share_1290939/shadowcun/evals_dev/{task}/samples/'\n",
    "# dir_prompts = f'/apdcephfs_cq3/share_1290939/shadowcun/evals_dev/{task}/prompts/'\n",
    "# prompt_paths = [os.path.join(dir_prompts, x) for x in os.listdir('/apdcephfs_cq3/share_1290939/shadowcun/evals_dev/'+task+'/prompts')]\n",
    "dir_prompts = f'/apdcephfs_cq3/share_1290939/raphaelliu/evals_dev/{task}/prompts_updated/'\n",
    "prompt_paths = [os.path.join(dir_prompts, x) for x in os.listdir('/apdcephfs_cq3/share_1290939/raphaelliu/evals_dev/'+task+'/prompts_updated')]\n",
    "\n",
    "video_paths = [os.path.join(dir_videos, os.path.splitext(os.path.basename(x))[0]+'.mp4') for x in prompt_paths]\n",
    "prompt_paths_object = []\n",
    "video_paths_object = []\n",
    "det_count = []\n",
    "keyword_object = []\n",
    "prompt_object = []\n",
    "det_count_mse_all = []\n",
    "det_num = 0\n",
    "object_count = 0\n",
    "pos_num = 0\n",
    "vid_num = 0\n",
    "last_video_name = ' '\n",
    "start_time = time.time()  # Start the timer\n",
    "logger.info(f\"Total video num {len(prompt_paths)}\")\n",
    "wandb.log({\"Total video num\": len(prompt_paths)\n",
    "            })\n",
    "\n",
    "for i in range(len(prompt_paths)):\n",
    "    with open(prompt_paths[i], \"r\") as f:\n",
    "        data = f.read() \n",
    "    for keyword in keywords:\n",
    "        num = len(re.findall(r'\\b' + re.escape(keyword) + r'(s|es)?\\b', data, re.IGNORECASE))      \n",
    "        if num > 0:\n",
    "            video_name = os.path.splitext(os.path.basename(prompt_paths[i]))[0]\n",
    "            if video_name != last_video_name:\n",
    "                vid_num +=1\n",
    "            last_video_name = video_name\n",
    "            io_args = {\n",
    "                'input_video': f'/apdcephfs_cq3/share_1290939/shadowcun/evals_dev/{task}/samples/{video_name}.mp4',\n",
    "                'output_mask_dir': f'./{task}/{timestamp}/{keyword}/masks/{video_name}', # save pred masks\n",
    "                'original_video': f'./{task}/{timestamp}/{keyword}/original_{video_name}.mp4', \n",
    "                'output_video': f'./{task}/{timestamp}/{keyword}/mask_{video_name}.mp4', # mask+frame vizualization, mp4 or avi, else the same as input video\n",
    "                'output_gif': f'./{task}/{timestamp}/{keyword}/gif_{video_name}.gif', # mask visualization\n",
    "            }\n",
    "            path = create_directories(io_args['output_mask_dir'])\n",
    "            \n",
    "            ## detection \n",
    "            grounding_caption = keyword\n",
    "            pred_list, det_count_frames = video_detection(io_args, segtracker_args, sam_args, aot_args, grounding_caption, box_threshold, text_threshold, box_size_threshold, reset_image)\n",
    "            det_num += 1 \n",
    "            det_count_avg = np.sum(det_count_frames) / len(det_count_frames)\n",
    "            det_count_mse = np.sum(np.abs(det_count_frames - det_count_avg) / max(np.max(det_count_frames - det_count_avg), 0.001)) / len(det_count_frames)\n",
    "            if det_count_avg:\n",
    "                pos_num += 1\n",
    "            det_count_mse_all.append(det_count_mse)\n",
    "            prompt_paths_object.append(prompt_paths[i])\n",
    "            keyword_object.append(keyword)\n",
    "            prompt_object.append(data)\n",
    "            acc  = pos_num / det_num\n",
    "            object_metric = 1 - sum(det_count_mse_all) / det_num\n",
    "            logging.info(f\"Det num: {det_num}, Positive num: {pos_num}, Acc: {acc}, Step obj-gen score: {det_count_mse}, Whole obj-gen score: {object_metric}, Vid num: {vid_num}, Prompt: {data}, Path index: {i}, Video name: {video_name}\")\n",
    "            wandb.log({\\\n",
    "                \"Positive rate\": acc,\n",
    "                \"Obj-gen score\":object_metric,\n",
    "                \"Det num\": det_num,\n",
    "                \"Positive num\": pos_num,\n",
    "            })\n",
    "            \n",
    "object_metric = pos_num/det_num\n",
    "# for i in range(len(obj_num)):\n",
    "#     object_count += abs(det_count[i] - obj_num[i])/(et_count[i] + obj_num[i]) \n",
    "# object_count_metric = 1 - (object_count/len(obj_num))  \n",
    "end_time = time.time()  # Stop the timer\n",
    "elapsed_time = end_time - start_time\n",
    "logger.info(f\"Object generation metric: {object_metric:.2f} \")\n",
    "# logger.info(f\"N objects generation metric: {object_count_metric:.2f} \")\n",
    "logger.info(f\"Total processing time: {elapsed_time:.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# object_metric = pos_num/det_num\n",
    "# # for i in range(len(obj_num)):\n",
    "# #     object_count += abs(det_count[i] - obj_num[i])/(et_count[i] + obj_num[i]) \n",
    "# # object_count_metric = 1 - (object_count/len(obj_num))  \n",
    "# # end_time = time.time()  # Stop the timer\n",
    "# # elapsed_time = end_time - start_time\n",
    "# object_metric\n",
    "\n",
    "np.sum(np.abs(det_count_frames - det_count_avg) / max(np.max(det_count_frames - det_count_avg), 0.001)) / len(det_count_frames)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set parameters for input and output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task = 'fulljourney_videos'\n",
    "# task = 'pikavideos'\n",
    "video_name = \"generated_video_ec706b3b-84c5-4169-b65d-e72efadf754f\"\n",
    "io_args = {\n",
    "    'input_video': f'/apdcephfs_cq3/share_1290939/shadowcun/evals_dev/{task}/samples/{video_name}.mp4',\n",
    "    'output_mask_dir': f'./{task}/masks/{video_name}', # save pred masks\n",
    "    'output_video': f'./{task}/{video_name}.mp4', # mask+frame vizualization, mp4 or avi, else the same as input video\n",
    "    'output_gif': f'./{task}/{video_name}.gif', # mask visualization\n",
    "}\n",
    "path = create_directories(io_args['output_mask_dir'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tuning Grounding-DINO and SAM on the First Frame for Good Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose good parameters in sam_args based on the first frame segmentation result\n",
    "# other arguments can be modified in model_args.py\n",
    "# note the object number limit is 255 by default, which requires < 10GB GPU memory with amp\n",
    "sam_args['generator_args'] = {\n",
    "        'points_per_side': 30,\n",
    "        'pred_iou_thresh': 0.8,\n",
    "        'stability_score_thresh': 0.9,\n",
    "        'crop_n_layers': 1,\n",
    "        'crop_n_points_downscale_factor': 2,\n",
    "        'min_mask_region_area': 200,\n",
    "    }\n",
    "\n",
    "# Set Text args\n",
    "'''\n",
    "parameter:\n",
    "    grounding_caption: Text prompt to detect objects in key-frames\n",
    "    box_threshold: threshold for box \n",
    "    text_threshold: threshold for label(text)\n",
    "    box_size_threshold: If the size ratio between the box and the frame is larger than the box_size_threshold, the box will be ignored. This is used to filter out large boxes.\n",
    "    reset_image: reset the image embeddings for SAM\n",
    "'''\n",
    "# grounding_caption = \"car.suv\"\n",
    "grounding_caption = \"tree\" #must have this class in the image, otherwise go wrong\n",
    "# grounding_caption = \"Obama\" \n",
    "box_threshold, text_threshold, box_size_threshold, reset_image = 0.6, 0.5, 0.5, True \n",
    "\n",
    "cap = cv2.VideoCapture(io_args['input_video'])\n",
    "frame_idx = 0\n",
    "segtracker = SegTracker(segtracker_args,sam_args,aot_args)\n",
    "segtracker.restart_tracker()\n",
    "with torch.cuda.amp.autocast():\n",
    "    while cap.isOpened():\n",
    "        ret, frame = cap.read()\n",
    "        frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
    "        count, pred_mask, annotated_frame = segtracker.detect_and_seg(frame, grounding_caption, box_threshold, \n",
    "                                                               text_threshold, box_size_threshold) # return a flag, whether detected or not\n",
    "        torch.cuda.empty_cache()\n",
    "        obj_ids = np.unique(pred_mask)\n",
    "        obj_ids = obj_ids[obj_ids!=0]\n",
    "        print(\"processed frame {}, obj_num {}\".format(frame_idx,len(obj_ids)),end='\\n')\n",
    "        break\n",
    "    cap.release()\n",
    "    \n",
    "    init_res = draw_mask(annotated_frame, pred_mask,id_countour=False)\n",
    "    plt.figure(figsize=(10,10))\n",
    "    plt.axis('off')\n",
    "    plt.imshow(init_res)\n",
    "    plt.show()\n",
    "    plt.figure(figsize=(10,10))\n",
    "    plt.axis('off')\n",
    "    plt.imshow(colorize_mask(pred_mask))\n",
    "    plt.show()\n",
    "\n",
    "    del segtracker\n",
    "    torch.cuda.empty_cache()\n",
    "    gc.collect()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Results for the Whole Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For every sam_gap frames, we use SAM to find new objects and add them for tracking\n",
    "# larger sam_gap is faster but may not spot new objects in time\n",
    "segtracker_args = {\n",
    "    'sam_gap': 49, # the interval to run sam to segment new objects\n",
    "    'min_area': 200, # minimal mask area to add a new mask as a new object\n",
    "    'max_obj_num': 255, # maximal object number to track in a video\n",
    "    'min_new_obj_iou': 0.8, # the area of a new object in the background should > 80% \n",
    "}\n",
    "\n",
    "# source video to segment\n",
    "cap = cv2.VideoCapture(io_args['input_video'])\n",
    "fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "# output masks\n",
    "output_dir = io_args['output_mask_dir']\n",
    "if not os.path.exists(output_dir):\n",
    "    os.makedirs(output_dir)\n",
    "pred_list = []\n",
    "masked_pred_list = []\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()\n",
    "sam_gap = segtracker_args['sam_gap']\n",
    "frame_idx = 0\n",
    "segtracker = SegTracker(segtracker_args, sam_args, aot_args)\n",
    "segtracker.restart_tracker()\n",
    "\n",
    "with torch.cuda.amp.autocast():\n",
    "    while cap.isOpened():\n",
    "        ret, frame = cap.read()\n",
    "        if not ret:\n",
    "            break\n",
    "        frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
    "        if frame_idx == 0:\n",
    "            pred_mask, _ = segtracker.detect_and_seg(frame, grounding_caption, box_threshold, text_threshold, box_size_threshold, reset_image)\n",
    "            # pred_mask = cv2.imread('./debug/first_frame_mask.png', 0)\n",
    "            torch. cuda.empty_cache()\n",
    "            gc.collect()\n",
    "            segtracker.add_reference(frame, pred_mask)\n",
    "        elif (frame_idx % sam_gap) == 0:\n",
    "            seg_mask, _ = segtracker.detect_and_seg(frame, grounding_caption, box_threshold, text_threshold, box_size_threshold, reset_image)\n",
    "            save_prediction(seg_mask, './debug/seg_result', str(frame_idx)+'.png')\n",
    "            torch.cuda.empty_cache()\n",
    "            gc.collect()\n",
    "            track_mask = segtracker.track(frame)\n",
    "            save_prediction(track_mask, './debug/aot_result', str(frame_idx)+'.png')\n",
    "            # find new objects, and update tracker with new objects\n",
    "            new_obj_mask = segtracker.find_new_objs(track_mask, seg_mask)\n",
    "            if np.sum(new_obj_mask > 0) >  frame.shape[0] * frame.shape[1] * 0.4:\n",
    "                new_obj_mask = np.zeros_like(new_obj_mask)\n",
    "            save_prediction(new_obj_mask,output_dir,str(frame_idx)+'_new.png')\n",
    "            pred_mask = track_mask + new_obj_mask\n",
    "            # segtracker.restart_tracker()\n",
    "            segtracker.add_reference(frame, pred_mask)\n",
    "        else:\n",
    "            pred_mask = segtracker.track(frame,update_memory=True)\n",
    "        torch.cuda.empty_cache()\n",
    "        gc.collect()\n",
    "        \n",
    "        save_prediction(pred_mask,output_dir,str(frame_idx)+'.png')\n",
    "        # masked_frame = draw_mask(frame,pred_mask)\n",
    "        # masked_pred_list.append(masked_frame)\n",
    "        # plt.imshow(masked_frame)\n",
    "        # plt.show() \n",
    "        \n",
    "        pred_list.append(pred_mask)\n",
    "        print(\"processed frame {}, obj_num {}\".format(frame_idx,segtracker.get_obj_num()),end='\\r')\n",
    "        frame_idx += 1\n",
    "    cap.release()\n",
    "    print('\\nfinished')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save results for visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# draw pred mask on frame and save as a video\n",
    "cap = cv2.VideoCapture(io_args['input_video'])\n",
    "fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "\n",
    "if io_args['input_video'][-3:]=='mp4':\n",
    "    fourcc =  cv2.VideoWriter_fourcc(*\"mp4v\")\n",
    "elif io_args['input_video'][-3:] == 'avi':\n",
    "    fourcc =  cv2.VideoWriter_fourcc(*\"MJPG\")\n",
    "    # fourcc = cv2.VideoWriter_fourcc(*\"XVID\")\n",
    "else:\n",
    "    fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))\n",
    "out = cv2.VideoWriter(io_args['output_video'], fourcc, fps, (width, height))\n",
    "\n",
    "frame_idx = 0\n",
    "while cap.isOpened():\n",
    "    ret, frame = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "    frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
    "    pred_mask = pred_list[frame_idx]\n",
    "    masked_frame = draw_mask(frame,pred_mask)\n",
    "    # masked_frame = masked_pred_list[frame_idx]\n",
    "    masked_frame = cv2.cvtColor(masked_frame,cv2.COLOR_RGB2BGR)\n",
    "    out.write(masked_frame)\n",
    "    print('frame {} writed'.format(frame_idx),end='\\r')\n",
    "    frame_idx += 1\n",
    "out.release()\n",
    "cap.release()\n",
    "print(\"\\n{} saved\".format(io_args['output_video']))\n",
    "print('\\nfinished')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save colorized masks as a gif\n",
    "duration = 1000 / fps\n",
    "imageio.mimsave(io_args['output_gif'],pred_list,duration=duration)\n",
    "print(\"{} saved\".format(io_args['output_gif']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# manually release memory (after cuda out of memory)\n",
    "del segtracker\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "vscode": {
   "interpreter": {
    "hash": "536611da043600e50719c9460971b5220bad26cd4a87e5994bfd4c9e9e5e7fb0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
